#!/bin/bash
# https://github.com/ContextualAI/gritlm commit_id：724df95c0f760249d3581f82cd7ca7f9ad5191c0
# 脚本丢在 run.py 的同级目录下，修改 run.py 的导包路径
# 修改config目录下对应的加速配置，fsdp_transformer_layer_cls_to_wrap: Qwen2DecoderLayer
# 用 accelerate 加速，并启动


export WANDB_DISABLED=true
# SBATCH --job-name=gritlm
# SBATCH --nodes=1
# SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
# SBATCH --hint=nomultithread         # we get physical cores not logical
# SBATCH --partition=a3
# SBATCH --gres=gpu:8                 # number of gpus
# SBATCH --time 999:00:00             # maximum execution time (HH:MM:SS)
# SBATCH --output=/data/niklas/jobs/%x-%j.out           # output file name
# SBATCH --exclusive

######################
### Set enviroment ###
######################
# cd /home/niklas/gritlm/gritlm
# source /env/bin/start-ctx-user
# conda activate gritlm
#NCCL_ASYNC_ERROR_HANDLING=1
# export WANDB_PROJECT="gritlm"
# Training setup
GPUS_PER_NODE=8
# so processes know who to talk to
MASTER_ADDR=127.0.0.1
MASTER_PORT=29500
NNODES=4
NODE_RANK=0
WORLD_SIZE=32
######################

######################
#### Set network #####
######################
# head_node_ip=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
######################


LAUNCHER="accelerate launch \
    --config_file /home/jovyan/pcj-backup/workdesk/gritlm/scripts/configs/config_32gpusfsdp_m7.yml \
    --num_machines $NNODES \
    --num_processes $WORLD_SIZE \
    --main_process_ip "$MASTER_ADDR" \
    --main_process_port $MASTER_PORT \
    --num_processes $WORLD_SIZE \
    --machine_rank  $NODE_RANK \
    --rdzv_conf rdzv_backend=c10d \
    --max_restarts 0 \
    --tee 3 \
    "

export CMD=" \
    -m run \
    --output_dir /home/jovyan/pcj-backup/workdesk/gritlm/output \
    --model_name_or_path /home/jovyan/pcj-backup/models/gte_Qwen2-7B-instruct/iic/gte_Qwen2-7B-instruct \
    --train_data xxx.jsonl \
    --learning_rate 2e-5 \
    --lr_scheduler_type linear \
    --warmup_ratio 0.03 \
    --max_steps 1000 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 1 \
    --dataloader_drop_last \
    --normalized \
    --temperature 0.02 \
    --train_group_size 1 \
    --query_max_len 128 \
    --passage_max_len 128 \
    --mode embedding \
    --logging_steps 10 \
    --bf16 \
    --pooling_method mean \
    --attn bbcc \
    --attn_implementation sdpa \
    --gradient_checkpointing \
    --save_steps 200
    "

$LAUNCHER $CMD
